package com.cloudea.learn.layer;

import com.cloudea.learn.Tensor;

import java.util.*;

/**
 * 组件的基类
 * @author cloudea
 * @date 2021/12/29
 */
public class Layer {
    public static final String NODE_ADDR_NAME = "node";
    private String name;

    public Layer(){
        this("layer");
    }

    public Layer(String name){
        this.name = String.valueOf(name);
    }

    /**
     * 前向输入多个Tensor
     * 拒绝过度设计
     * @param inputs 输入
     * @return 输出
     */
    public final Tensor call(Tensor... inputs){
        Node node = new Node();
        for (Tensor input:inputs){
            Node prev = (Node) input.getAttribute(NODE_ADDR_NAME);
            node.prev.add(prev);
            if(prev != null){
                prev.next.add(node);
            }
        }

        //前向计算
        Context context = new Context();
        Tensor output = forward(context, inputs);
        node.inputs = inputs;
        node.output = output;
        node.layer = this;
        node.context = context;

        //如果内部计算没有结点产生的话，则说明当前层需要计算梯度
        Node outputNode = (Node) output.getAttribute(NODE_ADDR_NAME);
        if(outputNode == null){
            output.setAttribute(NODE_ADDR_NAME, node);
        }
        return output;
    }

    /**
     * 前向输入多个Tensor
     * 拒绝过度设计
     * @param inputs 输入
     * @return 输出
     */
    protected Tensor forward(Context context, Tensor... inputs){
        return new Tensor(0);
    }

    /**
     * 梯度信息的反向传播
     * @param grid 损失函数对该单元输出的梯度
     * @return 损失函数对该单元输入的梯度,默认返回null
     */
    public Tensor[] backward(Context context, Tensor grid, Tensor... inputs){
        return new Tensor[]{grid};
    }


    /**
     * 计算结点，用于构建动态图
     */
    public static class Node implements Comparable<Node>{
        public List<Node> prev = new LinkedList<>();    //前向结点
        public List<Node> next = new LinkedList<>();    //后向结点
        public Layer layer ;                            // 对应的计算层,如果没有层与之对应，则说明是参数结点
        public Context context;
        public Tensor[] inputs;
        public Tensor output;                           // 如果是参数结点，则存储是的参数
        public List<Tensor> grids = new LinkedList<>(); // 目标函数对算子的梯度，因为算子的值可能被多个其它算子引用，因此需要收集全部的梯度才能求总的梯度

        @Override
        public int compareTo(Node node) {
            int nextSizeSpan = this.next.size() - node.next.size() ;
            if(nextSizeSpan > 0){
                return 1;
            } else if (nextSizeSpan < 0){
                return -1;
            }
            return 0;
        }
    }

    /**
     * 动态图中结点的上下文对象
     * 用于存储前向计算中的一些信息
     */
    public static class Context {
        private final Map<String, Object> attributes = new HashMap<>();

        public void setAttribute(String name, Object object) {
            this.attributes.put(name, object);
        }

        public Object getAttribute(String name) {
            return attributes.get(name);
        }
    }
}
